import os
import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import math

# 导入tkinterdnd2实现拖放
from tkinterdnd2 import DND_FILES, TkinterDnD

CANVAS_WIDTH, CANVAS_HEIGHT = 636, 900

#################################################################################
#名称： DraggableImage
#作用： 可拖拽的图片类，支持拖拽、缩放、旋转等操作
#参数： 无
#返回： 无
#注意： 继承此类需要传入canvas、PIL图片对象和初始坐标
#################################################################################
class DraggableImage:
    #################################################################################
    #名称： __init__
    #作用： 初始化可拖拽图片对象
    #参数： canvas-画布对象, pil_image-PIL图片对象, x-初始x坐标, y-初始y坐标
    #返回： 无
    #注意： 会自动在画布上创建图片并绑定事件
    #################################################################################
    def __init__(self, canvas, pil_image, x, y):
        self.canvas = canvas
        self.original_image = pil_image.convert("RGBA")
        self.x = x
        self.y = y
        self.scale = 1.0
        self.angle = 0
        self.drag_data = {"x": 0, "y": 0}

        self.tk_image = ImageTk.PhotoImage(self.original_image)
        self.image_id = canvas.create_image(x, y, image=self.tk_image)

        self.canvas.tag_bind(self.image_id, '<Button-1>', self.on_select)
        self.canvas.tag_bind(self.image_id, '<ButtonPress-1>', self.start_move)
        self.canvas.tag_bind(self.image_id, '<B1-Motion>', self.moving)

    #################################################################################
    #名称： on_select
    #作用： 处理图片被选中事件
    #参数： event-事件对象
    #返回： "break"-阻止事件传播
    #注意： 会将当前图片设为选中状态
    #################################################################################
    def on_select(self, event):
        self.canvas.master.selected_image = self
        return "break"

    #################################################################################
    #名称： start_move
    #作用： 开始拖拽移动，记录初始位置
    #参数： event-鼠标按下事件对象
    #返回： 无
    #注意： 记录拖拽起始坐标
    #################################################################################
    def start_move(self, event):
        self.drag_data["x"] = event.x
        self.drag_data["y"] = event.y

    #################################################################################
    #名称： moving
    #作用： 处理图片拖拽移动
    #参数： event-鼠标移动事件对象
    #返回： 无
    #注意： 实时更新图片位置和内部坐标
    #################################################################################
    def moving(self, event):
        dx = event.x - self.drag_data["x"]
        dy = event.y - self.drag_data["y"]
        self.canvas.move(self.image_id, dx, dy)
        self.x += dx
        self.y += dy
        self.drag_data["x"] = event.x
        self.drag_data["y"] = event.y

    #################################################################################
    #名称： update_image
    #作用： 根据当前缩放和旋转参数更新图片显示
    #参数： 无
    #返回： 无
    #注意： 会重新计算图片尺寸并更新画布显示
    #################################################################################
    def update_image(self):
        new_w = max(1, int(self.original_image.width * self.scale))
        new_h = max(1, int(self.original_image.height * self.scale))
        resized = self.original_image.resize((new_w, new_h), Image.LANCZOS)
        rotated = resized.rotate(self.angle, expand=True)
        self.tk_image = ImageTk.PhotoImage(rotated)
        self.canvas.itemconfig(self.image_id, image=self.tk_image)
        self.canvas.coords(self.image_id, self.x, self.y)

    #################################################################################
    #名称： render_to_canvas
    #作用： 将图片渲染到基础画布上并生成YOLO-OBB格式标注
    #参数： base_image-基础PIL图片对象
    #返回： YOLO-OBB格式的标注字符串
    #注意： 会将变换后的图片粘贴到基础图片上
    #################################################################################
    def render_to_canvas(self, base_image):
        new_w = max(1, int(self.original_image.width * self.scale))
        new_h = max(1, int(self.original_image.height * self.scale))
        resized = self.original_image.resize((new_w, new_h), Image.LANCZOS)
        rotated = resized.rotate(self.angle, expand=True)
        img_w, img_h = rotated.size
        left = int(self.x - img_w / 2)
        top = int(self.y - img_h / 2)
        base_image.paste(rotated, (left, top), rotated)

        # 计算原始矩形的四个角点（未旋转时）
        orig_w = new_w
        orig_h = new_h
        half_w = orig_w / 2
        half_h = orig_h / 2
        
        # 四个角点相对于中心的坐标
        corners = [
            (-half_w, -half_h),  # 左上
            (half_w, -half_h),   # 右上
            (half_w, half_h),    # 右下
            (-half_w, half_h)    # 左下
        ]
        
        # 应用旋转变换
        angle_rad = math.radians(self.angle)
        cos_a = math.cos(angle_rad)
        sin_a = math.sin(angle_rad)
        
        rotated_corners = []
        for x, y in corners:
            rx = x * cos_a - y * sin_a + self.x
            ry = x * sin_a + y * cos_a + self.y
            # 归一化到[0,1]
            rx_norm = rx / CANVAS_WIDTH
            ry_norm = ry / CANVAS_HEIGHT
            rotated_corners.extend([rx_norm, ry_norm])
        
        # YOLO-OBB格式: class_id x1 y1 x2 y2 x3 y3 x4 y4
        coords_str = ' '.join(f"{coord:.6f}" for coord in rotated_corners)
        return f"0 {coords_str}\n"

class ImageComposerApp(TkinterDnD.Tk):  # 注意这里继承TkinterDnD.Tk
    def __init__(self):
        super().__init__()
        self.title("图片合成编辑器")
        self.geometry(f"{CANVAS_WIDTH}x{CANVAS_HEIGHT}")
        self.configure(bg='gray')

        self.canvas = tk.Canvas(self, width=CANVAS_WIDTH, height=CANVAS_HEIGHT, bg="white")
        self.canvas.pack()

        self.images = []
        self.selected_image = None
        self.save_index = 0

        self.canvas.bind("<MouseWheel>", self.on_mousewheel)
        self.canvas.bind("<Button-4>", self.on_mousewheel)
        self.canvas.bind("<Button-5>", self.on_mousewheel)
        self.bind_all("<Control-s>", self.ctrl_s_save)

        btn = tk.Button(self, text="导入PNG图片", command=self.import_image)
        btn.pack(side=tk.BOTTOM, fill=tk.X)

        # 绑定拖放事件（支持多个文件）
        self.drop_target_register(DND_FILES)
        self.dnd_bind('<<Drop>>', self.drop_files)

    def import_image(self):
        paths = filedialog.askopenfilenames(filetypes=[("PNG图片", "*.png")])
        for path in paths:
            self.add_image(path)

    def drop_files(self, event):
        # event.data包含文件路径，多个用空格隔开（Windows），Linux可能是列表形式
        files = self.split_drop_files(event.data)
        # 去重处理
        unique_files = list(set(files))
        for f in unique_files:
            if f.lower().endswith('.png'):
                self.add_image(f)
        return 'break'

    def split_drop_files(self, data):
        # Windows下文件路径有空格，用 {} 包裹
        result = []
        temp = ""
        in_brace = False
        for c in data:
            if c == '{':
                in_brace = True
                temp = ""
            elif c == '}':
                in_brace = False
                result.append(temp)
            elif c == ' ' and not in_brace:
                if temp:
                    result.append(temp)
                    temp = ""
            else:
                temp += c
        if temp:
            result.append(temp)
        return result

    def add_image(self, path):
        try:
            pil_img = Image.open(path).convert("RGBA")
            # 限制图片数量防止崩溃
            if len(self.images) >= 20:
                messagebox.showwarning("警告", "图片数量过多，已达到上限20张！")
                return
        except Exception as e:
            messagebox.showerror("错误", f"打开图片失败: {e}")
            return
        center_x, center_y = CANVAS_WIDTH // 2, CANVAS_HEIGHT // 2
        img = DraggableImage(self.canvas, pil_img, center_x, center_y)
        self.images.append(img)

    def on_mousewheel(self, event):
        # 获取鼠标下的图片对象
        item = self.canvas.find_closest(event.x, event.y)[0]
        img = None
        for image in self.images:
            if image.image_id == item:
                img = image
                self.selected_image = img
                break
        
        if img is None:
            return

        if event.num == 4 or event.delta > 0:
            delta = 1
        elif event.num == 5 or event.delta < 0:
            delta = -1
        else:
            delta = 0

        shift_pressed = (event.state & 0x0001) != 0

        if shift_pressed:
            img.angle = (img.angle + delta * 5) % 360
        else:
            factor = 1.0 + 0.1 * delta
            img.scale = max(0.1, img.scale * factor)

        img.update_image()

    def ctrl_s_save(self, event=None):
        if not self.images:
            messagebox.showinfo("提示", "当前无任何图片可保存！")
            return

        base_img = Image.new("RGBA", (CANVAS_WIDTH, CANVAS_HEIGHT), (255, 255, 255, 255))

        yolo_lines = []
        for img in self.images:
            label_line = img.render_to_canvas(base_img)
            yolo_lines.append(label_line)

        os.makedirs("output", exist_ok=True)
        
        # 查找可用的文件名
        while True:
            save_name = f"{self.save_index:06d}"
            out_png = f"output/{save_name}.png"
            out_txt = f"output/{save_name}.txt"
            if not os.path.exists(out_png) and not os.path.exists(out_txt):
                break
            self.save_index += 1

        base_img.convert("RGB").save(out_png)
        with open(out_txt, "w") as f:
            f.writelines(yolo_lines)

        messagebox.showinfo("保存成功", f"图片和标注已保存为：\n{out_png}\n{out_txt}")

        self.save_index += 1

        for img in self.images:
            self.canvas.delete(img.image_id)
        self.images.clear()
        self.selected_image = None

    def run(self):
        self.mainloop()


if __name__ == "__main__":
    app = ImageComposerApp()
    app.run()
